/*
 * author      : 吴中勤 
 * date        : 2019/3/13
 * description : TODO
 */

package org.wzq.learn.infrastructure.lib.trace.threadpool;

import brave.Tracing;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.cloud.sleuth.SpanNamer;
import org.springframework.web.context.request.RequestContextHolder;
import org.wzq.learn.infrastructure.lib.trace.CustomThreadContext;

import java.util.concurrent.*;

/**
 * 自定义的线程池, 在此线程池中 使用 feignClient, 可以维持调用链不断
 */
public class CustomThreadPool {
    private static final Logger logger = LoggerFactory.getLogger(CustomThreadPool.class);

    private final int keepAliveTime = 120 * 60; // second
    private final ThreadPoolExecutor executorPool;
    private final ThreadExceptionHandler exceptionHandler;
    private final Tracing tracing;
    private final SpanNamer spanNamer;

    public CustomThreadPool(Tracing tracing, SpanNamer spanNamer) {
        this(tracing, spanNamer, 8, 16, 100);
    }

    public CustomThreadPool(Tracing tracing, SpanNamer spanNamer, int corePoolSize, int maxPoolSize, int queueSize) {
        this.tracing = tracing;
        this.spanNamer = spanNamer;
        this.exceptionHandler = new ThreadExceptionHandler();
        executorPool = new ThreadPoolExecutor(
                corePoolSize,
                maxPoolSize,
                keepAliveTime,
                TimeUnit.SECONDS,
                new ArrayBlockingQueue<Runnable>(queueSize)
        );
    }


    public void executor(Runnable r) {
        if (logger.isInfoEnabled()) logger.info("executorPool queue size=" + executorPool.getQueue().size());
        RunnableWrapper wrapper = new RunnableWrapper(
                r,
                CustomThreadContext.currentAttributes(),
                RequestContextHolder.currentRequestAttributes(),
                exceptionHandler,
                tracing,
                spanNamer);
        executorPool.execute(wrapper);
    }

    public <T> Future<T> submit(Callable<T> t) {
        CallableWrapper wrapper = new CallableWrapper<T>(
                t,
                CustomThreadContext.currentAttributes(),
                RequestContextHolder.currentRequestAttributes(),
                exceptionHandler,
                tracing,
                spanNamer);
        return (Future<T>) executorPool.submit(wrapper);
    }

}
